from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import List, Optional

mcp = FastMCP("Scikit-learn MCP Server")

"""
面向软件：python-scikit-learn

MCP tools 列表：
1. scikit_learn_check_installation - 检查 scikit-learn 是否已安装
   参数：无
   返回：统一 JSON 结构，包含 pip freeze 命令的执行结果

2. scikit_learn_run_test - 运行 scikit-learn 测试脚本并验证输出
   参数：test_script_path - 测试脚本路径，expected_result - 期望的结果字符串
   返回：统一 JSON 结构，包含测试脚本执行和结果验证的结果
"""

@mcp.tool()
def scikit_learn_check_installation() -> dict:
    """
    检查 scikit-learn 是否已正确安装
    
    用途：通过 pip3 freeze 命令检查 scikit-learn 包是否已安装
    参数：无
    返回：JSON 结构包含命令执行结果
        {
          "success": bool,
          "command": str,
          "exit_code": int,
          "stdout": str,
          "stderr": str
        }
    """
    command = "pip3 freeze | grep scikit-learn"
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout.strip(),
            "stderr": result.stderr.strip()
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def scikit_learn_run_test(test_script_path: str, expected_result: str) -> dict:
    """
    运行 scikit-learn 测试脚本并验证输出结果
    
    用途：执行指定的 Python 测试脚本，并使用 grep 验证输出是否包含期望的结果
    参数：
        test_script_path: 测试脚本的文件路径
        expected_result: 期望在输出中匹配的结果字符串
    返回：JSON 结构包含命令执行结果
        {
          "success": bool,
          "command": str,
          "exit_code": int,
          "stdout": str,
          "stderr": str
        }
    """
    command = f"python3 {test_script_path} | grep \"{expected_result}\""
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=60)
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout.strip(),
            "stderr": result.stderr.strip()
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()